# 从当前包导入基础大语言模型抽象基类
from .base import BaseLLM

# 导入OpenAI库，用于调用大模型API
from openai import OpenAI

# 从配置文件导入火山豆包API密钥、基础URL和模型名称
from config import DOUBAO_API_KEY, DOUBAO_BASE_URL, DOUBAO_MODEL


# 定义豆包大语言模型类，继承自BaseLLM
class DoubaoLLM(BaseLLM):
    # 初始化方法
    def __init__(self):
        # 创建OpenAI客户端，指定基础URL和API密钥
        self.client = OpenAI(base_url=DOUBAO_BASE_URL, api_key=DOUBAO_API_KEY)

    # 实现generate方法，用于生成模型回复
    def generate(self, prompt, **kwargs):
        # 调用OpenAI的chat.completions.create方法生成回复
        resp = self.client.chat.completions.create(
            # 指定使用的模型名称
            model=DOUBAO_MODEL,
            # 构造对话消息列表，用户输入为prompt
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": prompt,
                        },
                    ],
                }
            ],
            # 设置温度参数，默认为0.3
            temperature=kwargs.get("temperature", 0.3),
            # 设置最大token数，默认为4096
            max_tokens=kwargs.get("max_tokens", 4096),
        )
        # 返回模型生成的回复内容
        return resp.choices[0].message.content
